#ifndef _GCOOGRAPH_CUH__
#define _GCOOGRAPH_CUH__

#include "GraphData.cuh"
#include "config.hxx"

// #define DEBUG
#include "../cuda/cuda_utils.cuh"
#include "../utils.hxx"

namespace GraphData
{
  template <typename IndexType>
  struct GCOOGraph : public MetaGraph<IndexType>
  {
    const static int GRP_SIZE = 4;
    IndexType *src_list, *dst_list, *grp_offset; /// the new structure deprecated the combined storage
    IndexType *edge_idx;
    size_t ngrps;

    using super = MetaGraph<IndexType>;

    GCOOGraph( size_t nvertex, size_t nedge,
               IndexType *src_list, IndexType *dst_list, IndexType *edge_idx,
               IndexType *group_offset, size_t num_groups, DeviceTag dev = CPU
             ): MetaGraph<IndexType>(nvertex, nedge, dev), src_list(src_list), dst_list(dst_list), 
                grp_offset(group_offset), ngrps(num_groups)
    {
      super::independent = false;
    }

    GCOOGraph( IndexType *src_raw, IndexType *dst_raw, IndexType *eidx_raw, size_t nedge):
               MetaGraph<IndexType>(src_raw, dst_raw, eidx_raw, nedge)
    {
      _initializeField(super::nvertex, super::nedge);
      MYLOG("COO init Ok\n");
      fromCOO(src_raw, dst_raw, eidx_raw);
    }

    void fromCOO(IndexType *src_raw, IndexType *dst_raw, IndexType *eidx_raw)
    {
      using namespace std;
      struct edge_t
      {
        IndexType s, d;
        IndexType val;
        inline IndexType gid() const { return d / GRP_SIZE; }
      };
      
      bool with_eidx = (eidx_raw != nullptr);
      vector<edge_t> edge_aos;
      vector<IndexType> grp_sizes(ngrps, 0);

      for (IndexType i=0; i<super::nedge; ++i)
      {
        IndexType gid = dst_raw[i] / GRP_SIZE;
        IndexType ew = with_eidx ? eidx_raw[i] : i;
        edge_aos.push_back({src_raw[i], dst_raw[i], ew});
        ++grp_sizes[gid];
      }

      sort(edge_aos.begin(), edge_aos.end(),
        [](const edge_t &a, const edge_t &b) -> bool
        {
          if (a.gid() == b.gid())
            return a.s < b.s;
          else return a.gid() < b.gid();   
        }
      );
      MYLOG("Tag 2.1\n");
      /// prefix sum of grp_sizes
      vector<IndexType>grp_off(ngrps, 0);
      for (IndexType i=1; i<ngrps; ++i)
      {
        grp_off[i] = grp_sizes[i-1]+grp_off[i-1];
        grp_offset[i] = grp_off[i];
        MYLOG(" -- grp %d size = %d\n", i, grp_off[i]);
      }
      grp_offset[0] = 0;
      grp_offset[ngrps] = super::nedge;

      MYLOG("Tag 2.2\n");
      for (auto i=edge_aos.begin(); i<edge_aos.end(); ++i)
      {
        MYLOG(" -- puting %d->%d to grp %d @[%d]\n", i->s, i->d, i->gid(), grp_off[i->gid()]);
        src_list[grp_off[i->gid()]] = i->s;
        dst_list[grp_off[i->gid()]] = i->d;
        edge_idx[grp_off[i->gid()]] = i->val;
        ++grp_off[i->gid()]; 
      }
    }

    ~GCOOGraph()
    {
      if (super::independent)
      {
        MEMFREE(src_list, super::on_dev);
        MEMFREE(dst_list, super::on_dev);
        MEMFREE(grp_offset, super::on_dev);
        MEMFREE(edge_idx, super::on_dev);
      }
    }

    // void toTensor(std::vector<torch::Tensor>& tlist)
    // {
    //   ASSERT(super::on_dev == CPU, "device copy is not available yet");
    //   ASSERT(tlist.size() == 3, "expect, 3 tensor");
    //   super::independent = false;

    //   torch::ScalarType indTp, valTp;
    //   if (sizeof(IndexType) == 4) {
    //     indTp = torch::kInt32;
    //   } else if (sizeof(IndexType) == 8) {
    //     indTp = torch::kInt64;
    //   } else {
    //     ASSERT(false, "unsupported index type");
    //   }

    //   if (sizeof(IndexType) == 4) {
    //     valTp = torch::kFloat32;
    //   } else if (sizeof(IndexType) == 8) {
    //     valTp = torch::kFloat64;
    //   } else {
    //     ASSERT(false, "unsupported value type");
    //   }

    //   auto indOpt = torch::TensorOptions().dtype(indTp);
    //   auto valOpt = torch::TensorOptions().dtype(valTp);

    // #define PACK(tn, ptr, size, ty, opt) \
    //   tn = torch::zeros((size), opt);                     \
    //   memcpy(tn.data_ptr<ty>(), ptr, (size) * sizeof(ty));\
    //   free(ptr); ptr = tn.data_ptr<ty>()

    //   PACK(tlist[1], grp_offset, ngrps+1, IndexType, indOpt);
    //   if (super::weighted) {
    //     PACK(tlist[2], edge_idx, super::nedge, IndexType, valOpt);
    //   }
    // #undef PACK

    //   // the soa edge list is packed manually
    //   tlist[0] = torch::zeros(super::nedge*2, indOpt);
    //   auto edge_list_ptr = tlist[0].data_ptr<IndexType>();
    //   memcpy(edge_list_ptr, dst_list, super::nedge*sizeof(IndexType)); 
    //   memcpy(edge_list_ptr+super::nedge, src_list, super::nedge*sizeof(IndexType)); 
    //   free(src_list); free(dst_list);
    //   dst_list = edge_list_ptr;
    //   src_list = edge_list_ptr + super::nedge;
    
    // }

    void cpu()
    {
      if (super::on_dev != CPU)
      {

        IndexType *sl, *dl, *go;
        IndexType *ev;

        MEMALLOC(IndexType, sl, super::nedge, CPU);
        MEMALLOC(IndexType, dl, super::nedge, CPU);
        MEMALLOC(IndexType, go, ngrps+1, CPU);
        MEMALLOC(IndexType, ev, super::nedge, CPU);
        
        H_ERR(TOHOST(src_list, sl, super::nedge));
        H_ERR(TOHOST(dst_list, dl, super::nedge));
        H_ERR(TOHOST(grp_offset, go, ngrps+1));
        H_ERR(TOHOST(edge_idx, ev, super::nedge));

        MEMFREE(src_list, super::on_dev);
        MEMFREE(dst_list, super::on_dev);
        MEMFREE(grp_offset, super::on_dev);
        MEMFREE(edge_idx, super::on_dev);

        super::on_dev = CPU;

        src_list = sl;
        dst_list = dl;
        grp_offset = go;
        edge_idx = ev;
      }
    }

    void cuda()
    {
      if (super::on_dev != CUDA)
      {
        IndexType *sl, *dl, *go;
        IndexType *ev;

        MEMALLOC(IndexType, sl, super::nedge, CUDA);
        MEMALLOC(IndexType, dl, super::nedge, CUDA);
        MEMALLOC(IndexType, go, ngrps+1, CUDA);
        MEMALLOC(IndexType, ev, super::nedge, CUDA);
        
        H_ERR(TODEV(sl, src_list, super::nedge));
        H_ERR(TODEV(dl, dst_list, super::nedge));
        H_ERR(TODEV(go, grp_offset, ngrps+1));
        H_ERR(TODEV(ev, edge_idx, super::nedge));

        MEMFREE(src_list, super::on_dev);
        MEMFREE(dst_list, super::on_dev);
        MEMFREE(grp_offset, super::on_dev);
        MEMFREE(edge_idx, super::on_dev);
          
        super::on_dev = CUDA;

        src_list = sl;
        dst_list = dl;
        grp_offset = go;
        edge_idx = ev;
      }
    }

  private:
    void _initializeField( size_t nvertex, size_t nedge, DeviceTag dev = CPU )
    {
      super::independent = true;
      super::on_dev = dev;
      ngrps = CEIL(super::nvertex, GRP_SIZE);

      MEMALLOC(IndexType, src_list, super::nedge, dev);
      MEMALLOC(IndexType, dst_list, super::nedge, dev);
      MEMALLOC(IndexType, grp_offset, ngrps+1, dev);
      MEMALLOC(IndexType, edge_idx, super::nedge, dev)
    }
  };
};
#endif